from pydantic import BaseModel, Field
from langchain_core.prompts import ChatPromptTemplate
from patagent.llm import chat_llm
from patagent.prompt import RETRIEVAL_GRADE_PROMPT

class GradeDocuments(BaseModel):
    """Binary score for relevance check on retrieved documents."""

    binary_score: str = Field(
        description="Documents are relevant to the question, 'yes' or 'no'"
    )


# LLM with function call
structured_llm_grader = chat_llm.with_structured_output(GradeDocuments)

grade_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", RETRIEVAL_GRADE_PROMPT),
        ("human", "Retrieved document: \n\n {document} \n\n User question: {question}"),
    ]
)

retrieval_grader = grade_prompt | structured_llm_grader